# Copyright (c) Alibaba, Inc. and its affiliates.
import ast
import math
import os
from dataclasses import dataclass, field
from typing import Any, Dict, Literal, Optional, Union

import json
import torch
from transformers.utils import is_torch_mps_available

from swift.llm import MODEL_MAPPING, HfConfigFactory, get_model_info_meta, get_model_name
from swift.utils import get_dist_setting, get_logger

logger = get_logger()


@dataclass
class ModelArguments:
    """
    ModelArguments 类是一个数据类，包含与模型配置和使用相关的各种参数。

    参数：
    model（可选[str]）：model_id 或 model_path。默认值为 None。
    model_type（可选[str]）：模型组的类型。默认值为 None。
    model_revision（可选[str]）：模型的修订版本。默认值为 None。
    torch_dtype（字面量）：模型参数 dtype。默认值为 None。
    attn_impl（字面量）：要使用的注意实现。默认值为 None。
    num_labels（可选[int]）：分类任务的标签数量。默认值为 None。
    rope_scaling（字面量）：要使用的绳索缩放类型。默认值为 None。
    device_map（可选[str]）：设备映射配置。默认值为 None。
    local_repo_path（可选[str]）：模型的本地 GitHub 仓库路径。默认值为 None。
    """
    # 模型的路径或 ID（如 HuggingFace 上的模型名）
    model: Optional[str] = None

    # 模型类型，用于从 MODEL_MAPPING 中选择支持的模型结构
    model_type: Optional[str] = field(
        default=None, metadata={'help': f'model_type choices: {list(MODEL_MAPPING.keys())}'})

    # 模型的版本号或 Git 分支名
    model_revision: Optional[str] = None

    # 模型任务类型，如 causal_lm、分类或 embedding
    task_type: Literal['causal_lm', 'seq_cls', 'embedding'] = None

    # 模型精度类型：bfloat16、float16、float32
    torch_dtype: Literal['bfloat16', 'float16', 'float32', None] = None

    # 注意力实现方式：Flash Attention、SDPA（官方高性能实现）、Eager、Flex 等
    attn_impl: Literal['flash_attn', 'sdpa', 'eager', 'flex_attention', None] = None

    # 分类任务时的标签数量
    num_labels: Optional[int] = None

    # 问题类型：回归、多标签分类、单标签分类
    problem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] = None

    # RoPE 缩放方式（旋转位置编码）
    rope_scaling: Literal['linear', 'dynamic'] = None

    # 模型加载到设备的映射配置，例如 {'cuda:0': [0, 1]}
    device_map: Optional[Union[dict, str]] = None

    # 每个设备的最大内存配置，例如 {'cuda:0': '20GiB'}
    max_memory: Optional[Union[dict, str]] = None

    # 如果模型包含自定义代码，从 GitHub 下载的本地路径
    local_repo_path: Optional[str] = None

    @staticmethod
    def parse_to_dict(value: Union[str, Dict, None], strict: bool = True) -> Union[str, Dict]:
        """将 JSON 字符串或 JSON 文件解析为 Python 字典"""
        if value is None:
            value = {}
        elif isinstance(value, str):
            if os.path.exists(value):  # 文件路径
                with open(value, 'r', encoding='utf-8') as f:
                    value = json.load(f)
            else:  # 字符串形式的 JSON
                try:
                    value = json.loads(value)
                except json.JSONDecodeError:
                    if strict:
                        logger.error(f"无法解析字符串: '{value}'")
                        raise
        return value

    def _init_device_map(self):
        """初始化 device_map 配置，支持多卡分布式模式"""
        if self.device_map:
            self.device_map = self.parse_to_dict(self.device_map, strict=False)

        # 获取分布式训练设置（local_rank）
        _, local_rank, _, local_world_size = get_dist_setting()
        if local_world_size > 1 and isinstance(self.device_map, dict) and local_rank > 0:
            for k, v in self.device_map.items():
                if isinstance(v, int):
                    self.device_map[k] += local_rank

    def _init_max_memory(self):
        """初始化 max_memory 参数，适配多卡场景"""
        if isinstance(self.max_memory, str):
            try:
                self.max_memory = ast.literal_eval(self.max_memory)
            except Exception:
                pass
        self.max_memory = self.parse_to_dict(self.max_memory)
        _, local_rank, _, local_world_size = get_dist_setting()
        if local_world_size > 1 and isinstance(self.max_memory, dict) and local_rank > 0:
            for k in list(self.max_memory.keys()):
                if isinstance(k, int):
                    self.max_memory[k + local_rank] = self.max_memory.pop(k)

    def _init_torch_dtype(self) -> None:
        """根据训练类型/设备确定合适的 torch 数据类型"""
        from swift.llm import TrainArguments
        self.torch_dtype = HfConfigFactory.to_torch_dtype(self.torch_dtype)
        self.torch_dtype = self._init_model_info()

        # 如果是训练参数类的子类，初始化混合精度训练参数
        if isinstance(self, TrainArguments):
            self._init_mixed_precision()

    def _init_mixed_precision(self):
        """根据 torch_dtype 初始化混合精度（fp16/bf16）标志"""
        if is_torch_mps_available():
            fp16, bf16 = False, False
        elif self.torch_dtype in {torch.float16, torch.float32}:
            fp16, bf16 = True, False
        elif self.torch_dtype == torch.bfloat16:
            fp16, bf16 = False, True
        else:
            raise ValueError(f'args.torch_dtype: {self.torch_dtype}')
        if self.fp16 is None:
            self.fp16 = fp16
        if self.bf16 is None:
            self.bf16 = bf16

    def _init_rope_scaling(self):
        """根据模型最大长度初始化 rope_scaling 参数"""
        assert self.max_length is not None, '使用 rope_scaling 时必须设置 max_length'
        rope_scaling = self.model_info.rope_scaling or {}
        max_model_len = self.model_info.max_model_len
        rope_scaling_factor = 1.0
        if max_model_len:
            rope_scaling_factor = max(float(math.ceil(self.max_length / max_model_len)), 1.0)
        if rope_scaling:
            rope_scaling_factor = max(rope_scaling.get('factor', -1), rope_scaling_factor)
            rope_scaling['type'] = self.rope_scaling
            rope_scaling['factor'] = rope_scaling_factor
        else:
            rope_scaling = {'type': self.rope_scaling, 'factor': rope_scaling_factor}
        self.rope_scaling = rope_scaling
        logger.info(f'rope_scaling 配置为: {self.rope_scaling}')

    def _init_model_info(self) -> torch.dtype:
        """从模型配置中提取模型信息并初始化 task_type、num_labels 等"""
        self.model_info, self.model_meta = get_model_info_meta(**self.get_model_kwargs())
        self.task_type = self.model_info.task_type
        self.num_labels = self.model_info.num_labels
        self.model_dir = self.model_info.model_dir
        self.model_type = self.model_info.model_type

        # 初始化 rope_scaling
        if isinstance(self.rope_scaling, str):
            self._init_rope_scaling()

        return self.model_info.torch_dtype

    def __post_init__(self):
        """类初始化后执行的函数"""
        if self.model is None:
            raise ValueError(f'请设置 --model <model_id_or_path>，当前 model: {self.model}')
        self.model_suffix = get_model_name(self.model)
        self._init_device_map()
        self._init_max_memory()
        self._init_torch_dtype()

    def get_model_kwargs(self):
        """构造模型加载所需的关键参数字典"""
        return {
            'model_id_or_path': self.model,
            'torch_dtype': self.torch_dtype,
            'model_type': self.model_type,
            'revision': self.model_revision,
            'use_hf': self.use_hf,
            'hub_token': self.hub_token,
            'local_repo_path': self.local_repo_path,
            'device_map': self.device_map,
            'max_memory': self.max_memory,
            'quantization_config': self.get_quantization_config(),
            'attn_impl': self.attn_impl,
            'rope_scaling': self.rope_scaling,
            'task_type': self.task_type,
            'num_labels': self.num_labels,
            'problem_type': self.problem_type,
        }
